package mr

import (
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net/rpc"
	"os"
	"sort"
	"strconv"
	"time"
)

// Map functions return a slice of KeyValue.
type KeyValue struct {
	Key   string
	Value string
}

// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

// keyReduceIndex(key, nReduce), according to nReduce to decide which reduce to do the work
func keyReduceIndex(key string, nReduce int) int {
	return ihash(key) % nReduce
}

func makeIntermediateFromFile(filename string, mapf func(string, string) []KeyValue) []KeyValue {
	path := filename

	file, err := os.Open(path)

	if err != nil {
		log.Fatalf("cannot open %v", path)
	}
	// readfile content
	content, err := io.ReadAll(file)
	if err != nil {
		log.Fatalf("cannot read file %v", filename)
	}
	file.Close()
	// generate k/v array
	kva := mapf(filename, string(content))

	return kva
}

func askTask() *TaskReqReply {
	log.Println("Worker: start asking for a task")
	reply := &TaskReqReply{}
	call("Coordinator.TaskRequest", &TaskReqArg{}, reply)
	return reply
}

func doMapTask(mapf func(string, string) []KeyValue, reply *TaskReqReply) {
	log.Printf("Worker: start doing map task on file %v %v\n", reply.Id, reply.FileName)
	intermediate := makeIntermediateFromFile(reply.FileName, mapf)
	// hash into bucket
	buckets := make([][]KeyValue, reply.NReduce)
	for i := range buckets {
		buckets[i] = []KeyValue{}
	}
	for _, kva := range intermediate {
		hash := keyReduceIndex(kva.Key, reply.NReduce)
		buckets[hash] = append(buckets[hash], kva)
	}
	// write into intermedite file
	for i := range buckets {
		fname := "mr-" + strconv.Itoa(reply.Id) + "-" + strconv.Itoa(i)
		log.Printf("start write file %v", fname)
		ofile, _ := os.CreateTemp("/tmp", fname+"*")
		if err := json.NewEncoder(ofile).Encode(buckets[i]); err != nil {
			log.Fatalf("cannot write intermediate data to %v\n", ofile.Name())
		}

		if err := os.Rename(ofile.Name(), fname); err != nil {
			log.Fatalf("rename file error: %v", err)
		}
		ofile.Close()
	}
	call("Coordinator.TaskReport", &ResponseArg{Id: reply.Id, FileName: reply.FileName, Type: reply.Type}, &ResponseReply{})
}

func doReduceTask(reducef func(string, []string) string, reply *TaskReqReply) {
	log.Println("Worker: start doing reduce task")
	var intermediate []KeyValue
	for i := 0; i < reply.NMap; i++ {
		fname := "mr-" + strconv.Itoa(i) + "-" + strconv.Itoa(reply.Id)
		data, err := os.ReadFile(fname)
		if err != nil {
			log.Fatalf("Read intermrdiate file %v failed! err: %v\n", fname, err)
		}
		var temp []KeyValue
		if err := json.Unmarshal(data, &temp); err != nil {
			log.Fatal("Unmarshal data failed!")
		}
		intermediate = append(intermediate, temp...)
	}
	sort.Sort(ByKey(intermediate))

	oname := "mr-out-" + strconv.Itoa(reply.Id)
	ofile, _ := os.Create(oname)
	i := 0
	for i < len(intermediate) {
		j := i + 1
		for j < len(intermediate) && intermediate[j].Key == intermediate[i].Key {
			j++
		}
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, intermediate[k].Value)
		}
		output := reducef(intermediate[i].Key, values)

		// this is the correct format for each line of Reduce output.
		fmt.Fprintf(ofile, "%v %v\n", intermediate[i].Key, output)

		i = j
	}
	ofile.Close()
	call("Coordinator.TaskReport", &ResponseArg{Id: reply.Id, FileName: reply.FileName, Type: reply.Type}, &ResponseReply{})
}

// main/mrworker.go calls this function.
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.
	// Init the work info
	// uncomment to send the Example RPC to the coordinator.
	// CallExample()
	for {
		reply := askTask()
		log.Printf("Worker: receive coordinator's task %v\n", reply)
		switch reply.Type {
		case MapTask:
			doMapTask(mapf, reply)
		case ReduceTask:
			doReduceTask(reducef, reply)
		case WaitTask:
			time.Sleep(1 * time.Second)
		case ExistTask:
			log.Println("Worker: all task done, exiting")
		default:
			panic(fmt.Sprintf("Unkown taskType: %v", reply.Type))
		}
	}

}

// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	// the "Coordinator.Example" tells the
	// receiving server that we'd like to call
	// the Example() method of struct Coordinator.
	ok := call("Coordinator.Example", &args, &reply)
	if ok {
		// reply.Y should be 100.
		fmt.Printf("reply.Y %v\n", reply.Y)
	} else {
		fmt.Printf("call failed!\n")
	}
}

// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
